// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

namespace MQTTnet.Internal;

public sealed class AsyncSignal : IDisposable
{
    readonly object _syncRoot = new();

    bool _isDisposed;
    bool _isSignaled;
    AsyncSignalWaiter _waiter;

    public void Dispose()
    {
        lock (_syncRoot)
        {
            _waiter?.Dispose();
            _waiter = null;

            _isDisposed = true;
        }
    }

    public void Set()
    {
        lock (_syncRoot)
        {
            _isSignaled = true;

            Cleanup();

            // If there is already a waiting task let it run.
            if (_waiter != null)
            {
                _waiter.Approve();
                _waiter.Dispose();
                _waiter = null;

                // Since we already got a waiter the signal must be reset right now!
                _isSignaled = false;
            }
        }
    }

    public Task WaitAsync(CancellationToken cancellationToken = default)
    {
        cancellationToken.ThrowIfCancellationRequested();

        lock (_syncRoot)
        {
            ThrowIfDisposed();

            Cleanup();

            if (_isSignaled)
            {
                _isSignaled = false;
                return CompletedTask.Instance;
            }

            if (_waiter != null)
            {
                if (!_waiter.Task.IsCompleted)
                {
                    throw new InvalidOperationException("Only one waiting task is permitted per async signal.");
                }

                _waiter.Dispose();
            }

            _waiter = new AsyncSignalWaiter(cancellationToken);
            return _waiter.Task;
        }
    }

    void Cleanup()
    {
        // Cleanup if the previous waiter was cancelled.
        if (_waiter != null && _waiter.Task.IsCanceled)
        {
            _waiter = null;
        }
    }

    void ThrowIfDisposed()
    {
        ObjectDisposedException.ThrowIf(_isDisposed, this);
    }

    sealed class AsyncSignalWaiter : IDisposable
    {
        readonly AsyncTaskCompletionSource<bool> _promise = new();

        // ReSharper disable once FieldCanBeMadeReadOnly.Local
        CancellationTokenRegistration _cancellationTokenRegistration;

        volatile bool _isCompleted;

        public AsyncSignalWaiter(CancellationToken cancellationToken)
        {
            if (cancellationToken.CanBeCanceled)
            {
                _cancellationTokenRegistration = cancellationToken.Register(Cancel);
            }
        }

        public Task Task => _promise.Task;

        public void Approve()
        {
            if (_isCompleted)
            {
                return;
            }

            _isCompleted = true;
            _promise.TrySetResult(true);
        }

        public void Dispose()
        {
            _cancellationTokenRegistration.Dispose();

            if (_isCompleted)
            {
                // Avoid allocation of _ObjectDisposedException_ which may not be used.
                return;
            }

            _isCompleted = true;
            _promise.TrySetException(new ObjectDisposedException(nameof(AsyncSignalWaiter)));
        }

        void Cancel()
        {
            if (_isCompleted)
            {
                return;
            }

            _isCompleted = true;
            _promise.TrySetCanceled();
        }
    }
}